package com.yth.utils;

import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.ToString;

import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Parser for extracting structured blocks from LLM output text.
 *
 * @author yutianhong
 * @version 1.0
 * @since 2025/6/18 12:57
 */
public class AssistantMessageParser {
    /**
     * Type identifier for untagged text content.
     */
    private static final String TEXT_BLOCK_TYPE = "$$plain_text";

    // 正则表达式用于匹配标签
    private static final Pattern SINGLE_LINE_TAG_PATTERN = Pattern.compile("^\\s*<(\\w+)>(.*?)</\\1>\\s*$");
    private static final Pattern START_TAG_PATTERN = Pattern.compile("^\\s*<(\\w+)>(.*)$");
    private static final Pattern END_TAG_PATTERN = Pattern.compile("^(.*)</(\\w+)>\\s*$");

    private final Set<String> allowedTags;

    // Private constructor - use factory methods
    private AssistantMessageParser(Set<String> allowedTags) {
        if (allowedTags != null && allowedTags.isEmpty()) {
            throw new IllegalArgumentException("allowedTags cannot be empty");
        }
        this.allowedTags = allowedTags == null ? null
                : Collections.unmodifiableSet(new HashSet<>(allowedTags));
    }

    /**
     * Creates a parser that accepts all tag names.
     */
    public static AssistantMessageParser create() {
        return new AssistantMessageParser(null);
    }

    /**
     * Creates a parser that only accepts the specified tag names.
     */
    public static AssistantMessageParser create(Set<String> allowedTags) {
        return new AssistantMessageParser(allowedTags);
    }

    /**
     * Creates a streaming parser for line-by-line processing.
     */
    public StreamingParser createStreamingParser() {
        return new StreamingParser();
    }

    /**
     * Parses the entire input text and returns all extracted blocks.
     */
    public List<ContentBlock> parseAll(String input) {
        StreamingParser streamingParser = createStreamingParser();

        String[] lines = input.split("\n");
        for (String line : lines) {
            streamingParser.processLine(line);
        }
        // 处理可能未完成的标签
        streamingParser.finishPending();
        return streamingParser.currentBlocks();
    }

    private boolean isTagAllowed(String tagName) {
        return allowedTags == null || allowedTags.contains(tagName);
    }

    /**
     * Streaming parser for processing LLM output line by line.
     */
    public final class StreamingParser {
        private final List<ContentBlock> contentBlocks = new ArrayList<>();
        private final ParsingState state = new ParsingState();
        private final StringBuilder textBlockBuffer = new StringBuilder();

        /**
         * Processes a single line of input.
         */
        public void processLine(String line) {
            if (line == null) {
                return;
            }

            processLineInternal(line);
        }

        /**
         * Finalizes parsing and returns all extracted blocks.
         */
        public void finishPending() {
            if (state.isInTag()) {
                ContentBlock contentBlock = state.finishCurrentBlock();
                contentBlocks.add(contentBlock);
            } else {
                addTextBlockIfNonEmpty();
            }
        }

        /**
         * Returns a copy of currently parsed blocks without finalizing.
         */
        public List<ContentBlock> currentBlocks() {
            return Collections.unmodifiableList(new ArrayList<>(contentBlocks));
        }

        // Core parsing logic
        private void processLineInternal(String line) {
            if (line.trim().contains("\n")) {
                throw new IllegalArgumentException("Line contains newline character");
            }

            // Handle empty lines gracefully
            if (line.trim().isEmpty()) {
                if (state.isInTag()) {
                    state.appendContentLine("");
                }
                return;
            }

            // Try single-line complete tag first
            Matcher singleLineMatcher = SINGLE_LINE_TAG_PATTERN.matcher(line);
            if (singleLineMatcher.matches()) {
                String tagName = singleLineMatcher.group(1);
                String content = singleLineMatcher.group(2);
                if (isTagAllowed(tagName)) {
                    // Finish any pending block first
                    if (state.isInTag()) {
                        ContentBlock contentBlock = state.finishCurrentBlock();
                        contentBlocks.add(contentBlock);
                    } else {
                        addTextBlockIfNonEmpty();
                    }

                    contentBlocks.add(new ContentBlock(tagName, content.trim()));
                    return;
                }
            }

            // Check for end tag (with possible content before it)
            Matcher endTagMatcher = END_TAG_PATTERN.matcher(line);
            if (endTagMatcher.matches()) {
                String precedingContent = endTagMatcher.group(1);
                String tagName = endTagMatcher.group(2);

                if (state.isInTag() && tagName.equals(state.currentTagName())) {
                    if (!precedingContent.trim().isEmpty()) {
                        state.appendContentLine(precedingContent);
                    }
                    ContentBlock contentBlock = state.finishCurrentBlock();
                    contentBlocks.add(contentBlock);
                    return;
                }
            }

            // Check for start tag
            Matcher startTagMatcher = START_TAG_PATTERN.matcher(line);
            if (startTagMatcher.matches()) {
                String tagName = startTagMatcher.group(1);
                String trailingContent = startTagMatcher.group(2);

                if (isTagAllowed(tagName)) {
                    // Finish any pending block first
                    if (state.isInTag()) {
                        ContentBlock contentBlock = state.finishCurrentBlock();
                        contentBlocks.add(contentBlock);
                    } else {
                        addTextBlockIfNonEmpty();
                    }

                    state.startNewBlock(tagName);
                    if (!trailingContent.trim().isEmpty()) {
                        state.appendContentLine(trailingContent);
                    }
                    return;
                }
            }

            // append the entire line as content
            if (state.isInTag()) {
                state.appendContentLine(line);
            } else {
                appendTextBlockContentLine(line);
            }
        }

        private void appendTextBlockContentLine(String line) {
            if (textBlockBuffer.length() > 0) {
                textBlockBuffer.append('\n');
            }
            textBlockBuffer.append(line);
        }

        private void addTextBlockIfNonEmpty() {
            String text = textBlockBuffer.toString().trim();
            if (!text.isEmpty()) {
                contentBlocks.add(new ContentBlock(TEXT_BLOCK_TYPE, text));
                textBlockBuffer.setLength(0);
            }
        }
    }

    /**
     * Represents a parsed block with tag name and content.
     */
    @Getter
    @ToString
    @EqualsAndHashCode
    public static final class ContentBlock {
        private final String type;
        private final String content;

        ContentBlock(String type, String content) {
            this.type = Objects.requireNonNull(type, "type cannot be null");
            this.content = Objects.requireNonNull(content, "content cannot be null");
        }

        public boolean isTextBlock() {
            return TEXT_BLOCK_TYPE.equals(type);
        }
    }

    // Internal parsing state
    private static final class ParsingState {
        private String currentTagName;
        private final StringBuilder contentBuffer = new StringBuilder();

        boolean isInTag() {
            return currentTagName != null;
        }

        String currentTagName() {
            return currentTagName;
        }

        void startNewBlock(String tagName) {
            if (isInTag()) {
                throw new IllegalStateException(String.format("Already processing tag '%s', cannot start '%s'", currentTagName, tagName));
            }
            this.currentTagName = tagName;
            this.contentBuffer.setLength(0);
        }

        void appendContentLine(String contentLine) {
            if (!isInTag()) {
                throw new IllegalStateException("Not currently processing any tag");
            }
            if (contentBuffer.length() > 0) {
                contentBuffer.append('\n');
            }
            contentBuffer.append(contentLine);
        }

        ContentBlock finishCurrentBlock() {
            if (!isInTag()) {
                throw new IllegalStateException("Not currently processing any tag");
            }
            String content = contentBuffer.toString().trim();
            ContentBlock contentBlock = new ContentBlock(currentTagName, content);
            reset();
            return contentBlock;
        }

        private void reset() {
            currentTagName = null;
            contentBuffer.setLength(0);
        }
    }
}
